# encoding:utf-8
# -----------------------------------------------------------
# "ML-Stealer: Stealing Prediction Functionality of Machine Learning Models with Mere Black-Box Access"
# @author: Shijie Wang, 2019.
# ------------------------------------------------------------

import numpy as np
import pandas as pd
import torch
import statistics

from sklearn.externals import joblib
from synthetic import preprocess
from wgangp_adult import Generator
import scipy.stats


def modeltest(test_csv, fake_model, real_model):
    # get the test data
    test_data = pd.read_csv(test_csv, sep=',')
    test_data = preprocess(test_data.values)

    real_predict = real_model.predict_proba(test_data)
    fake_predict = fake_model(torch.tensor(test_data).float())

    return real_predict, fake_predict


def cross_entropy(targets, predictions, epsilon=1e-5):
    """
    Computes cross entropy between targets (encoded as one-hot vectors)
    and predictions.
    Input: predictions (N, k) ndarray
           targets (N, k) ndarray
    Returns: scalar
    """
    predictions = np.clip(predictions, epsilon, 1. - epsilon)
    N = predictions.shape[0]
    ce = -np.sum(targets*np.log(predictions+1e-9))/N
    return ce


def KL_CE_calculate(real_predict, fake_predict):
    kl_value = []
    ce_value = []
    for i in range(real_predict.shape[0]):
        kl_value.append(scipy.stats.entropy(fake_predict[i], real_predict[i]))
        ce_value.append(cross_entropy(real_predict[i], fake_predict[i]))
    return statistics.mean(kl_value), statistics.mean(ce_value)


def MSE_calculate(real_predict, fake_predict):
    MSE_Loss = ((np.max(real_predict, axis=1) - fake_predict[np.arange(len(real_predict)), np.argmax(real_predict, axis=1)]) ** 2).sum()
    return MSE_Loss / len(real_predict)


def Similarity_calculate(real_predict, fake_predict):
    error_num = 0
    for i in range(len(real_predict)):
        if fake_predict[i, np.argmax(real_predict, axis=1)[i]] < 0.5:
            error_num += 1
    ErrorRate = error_num / len(real_predict)
    Simialarity = (1 - ErrorRate) * 100
    return Simialarity


if __name__ == '__main__':
    adult_model = torch.load('./wgan_adult/yourpath', map_location=lambda storage, loc: storage)
    adult_rf_model = joblib.load('./victim_models/svm_part_adult.m')
    adult_model.eval()
    result_list = []

    for i in range(1, 21):
        real_predict, fake_predict = modeltest('./test_data/adult_test_{}.csv'.format(i), adult_model, adult_rf_model)
        MSE = MSE_calculate(real_predict, fake_predict.detach().numpy())
        Similarity = Similarity_calculate(real_predict, fake_predict.detach().numpy())
        result_list.append(np.hstack((MSE, Similarity)))
        print('MSE:{}, Similarity:{}'.format(MSE, Similarity))
    df_result = pd.DataFrame(result_list, columns=['MSE', 'Similarity'])
